#! /usr/bin/env python3

import argparse
from glob import glob
import os
import f90nml
import sys
sys.path.append(f'{os.path.dirname(os.path.realpath(__file__))}/../utils')
from utils import cli, check_files, run, parse_config, submit_job

def run_wrfda_update_bc(work_root, wrfda_root, config, args, update_lowbc=False, wrf_work_dir=None, wrfbdy=None, dom='d01', tag=None):
	start_time = config['custom']['start_time']
	datetime_fmt = 'YYYY-MM-DD_HH:mm:ss'
	start_time_str = start_time.format(datetime_fmt)
	max_dom = config['domains']['max_dom']

	if not wrf_work_dir:
		if tag != None:
			wrf_work_dir = f'{work_root}/{tag}/wrf'
		else:
			wrf_work_dir = f'{work_root}/wrf'

	dom_str = dom
	if tag != None:
		wrfda_work_dir = f'{work_root}/{tag}/wrfda/{dom_str}'
	else:
		wrfda_work_dir = f'{work_root}/wrfda/{dom_str}'
	if not os.path.isdir(wrfda_work_dir): os.mkdir(wrfda_work_dir)
	old_dir = os.getcwd()
	os.chdir(wrfda_work_dir)

	if not wrfbdy: wrfbdy = f'{wrf_work_dir}/wrfbdy_{dom_str}'

	cli.stage(f'Run WRFDA update_bc at {wrfda_work_dir} ...')

	expected_files = [wrfbdy, f'wrfvar_output_{start_time_str}', 'fg']
	if not check_files(expected_files):
		print(expected_files)
		cli.error('run_wrfda_update_bc: da_wrfvar.exe or real.exe wasn\'t executed successfully!')
	run(f'ln -sf {wrfbdy} wrfbdy_{dom_str}')
	run(f'ln -sf wrfvar_output_{start_time_str} wrfvar_output')

	parame_in = f90nml.read(f'{wrfda_root}/var/test/update_bc/parame.in')
	parame_in['control_param']['wrf_input'] = './fg'
	if update_lowbc:
		cli.notice('Update only low boundary condition.')
		parame_in['control_param']['low_bdy_only'] = True
	parame_in.write(f'{wrfda_work_dir}/parame.in', force=True)

	if update_lowbc:
		expected_file = f'wrfbdy_{dom_str}_{start_time_str}.low_updated'
	else:
		expected_file = f'wrfbdy_{dom_str}_{start_time_str}.lateral_updated'
	if not check_files(expected_file) or args.force:
		submit_job(f'{wrfda_root}/var/build/da_update_bc.exe', 1, config, args, wait=True)
		run(f'cp wrfbdy_{dom_str} {expected_file}')
	else:
		run(f'ls -l {expected_file}')

	cli.notice('Succeeded.')
	os.chdir(old_dir)
